package akka.agent.application

import akka.agent.domain.SessionEvent
import akka.javasdk.annotations.ComponentId
import akka.javasdk.annotations.Consume
import akka.javasdk.annotations.Query
import akka.javasdk.view.TableUpdater
import akka.javasdk.view.View
import java.util.Date

@ComponentId("view_chat_log")
class ConversationHistoryView : View() {
  data class ConversationHistory(val sessions: List<Session>)

  @JvmRecord data class Message(val message: String, val origin: String, val timestamp: Long)

  @JvmRecord
  data class Session(
      val userId: String,
      val sessionId: String,
      val creationDate: Long,
      val messages: List<Message>
  ) {
    fun add(message: Message) = copy(messages = messages + message)
  }

  @Query(
      "SELECT collect(*) as sessions FROM view_chat_log " +
          "WHERE userId = :userId ORDER by creationDate DESC")
  fun getSessionsByUser(userId: String): QueryEffect<ConversationHistory> = queryResult()
}

@Consume.FromEventSourcedEntity(SessionEntity::class)
class ChatMessageUpdater : TableUpdater<ConversationHistoryView.Session>() {
  fun onEvent(event: SessionEvent): Effect<ConversationHistoryView.Session> =
      when (event) {
        is SessionEvent.AiMessageAdded -> aiMessage(event)
        is SessionEvent.UserMessageAdded -> userMessage(event)
      }

  private fun aiMessage(
      event: SessionEvent.AiMessageAdded
  ): Effect<ConversationHistoryView.Session> {
    val newMessage =
        ConversationHistoryView.Message(
            message = event.response, origin = "ai", timestamp = event.fireDateTime.time)
    val rowState = rowStateOrNew(event.userId, event.sessionId)
    return effects().updateRow(rowState.add(newMessage))
  }

  private fun userMessage(
      event: SessionEvent.UserMessageAdded
  ): Effect<ConversationHistoryView.Session> {
    val newMessage =
        ConversationHistoryView.Message(
            message = event.query, origin = "user", timestamp = event.fireDateTime.time)
    val rowState = rowStateOrNew(userId = event.userId, sessionId = event.sessionId)
    return effects().updateRow(rowState.add(newMessage))
  }

  private fun rowStateOrNew(userId: String, sessionId: String): ConversationHistoryView.Session =
      if (rowState() != null) rowState()
      else
          ConversationHistoryView.Session(
              userId = userId,
              sessionId = sessionId,
              creationDate = Date().time,
              messages = listOf())
}
